import os
import rasterio
import numpy as np
import pandas as pd
# from rasterio.merge import merge
from tqdm import tqdm

# 行政编码到中文名的映射
admin_code_to_name_default = {
    "110000": "北京市",
    "120000": "天津市",
    "130000": "河北省",
    "140000": "山西省",
    "150000": "内蒙古自治区",
    "210000": "辽宁省",
    "220000": "吉林省",
    "230000": "黑龙江省",
    "310000": "上海市",
    "320000": "江苏省",
    "330000": "浙江省",
    "340000": "安徽省",
    "350000": "福建省",
    "360000": "江西省",
    "370000": "山东省",
    "410000": "河南省",
    "420000": "湖北省",
    "430000": "湖南省",
    "440000": "广东省",
    "450000": "广西壮族自治区",
    "460000": "海南省",
    "500000": "重庆市",
    "510000": "四川省",
    "520000": "贵州省",
    "530000": "云南省",
    "540000": "西藏自治区",
    "610000": "陕西省",
    "620000": "甘肃省",
    "630000": "青海省",
    "640000": "宁夏回族自治区",
    "650000": "新疆维吾尔自治区",
    "660000": "台湾省",
}

def analyze_tif_files(folder_path: str, admin_code_to_name=None, output_path: str = None) -> None:
    """
    分析指定文件夹中的 TIF 文件，并将结果保存到 Excel 文件中。

    :param folder_path: 包含 TIF 文件的文件夹路径
    :param admin_code_to_name: 行政编码到中文名的映射字典
    :param output_path: 结果保存的 Excel 文件路径，默认为 None，将根据 folder_path 生成
    """
    # 初始化一个空的DataFrame来存储结果
    if admin_code_to_name is None:
        admin_code_to_name = admin_code_to_name_default
    # 定义列名
    columns = ['行政编码', '核算地区', '实物量', '生态系统', '总值', '平均值']
    results = pd.DataFrame(columns=columns)

    # 遍历文件夹中的所有文件
    for filename in tqdm(os.listdir(folder_path), desc="Processing TIF files"):
        if filename.endswith('.tif'):
            # 提取行政编码
            admin_code = filename.split('_')[0]

            # 提取实物量类型和生态系统类型
            quantity_type = filename.split('_')[1]
            ecosystem_type = filename.split('_')[2]

            # 打开TIF文件
            with rasterio.open(os.path.join(folder_path, filename)) as src:
                # 读取所有波段的数据
                data = src.read(1)

                # 计算总值和平均值
                total_value = np.nansum(data)  # 忽略NaN值
                mean_value = np.nanmean(data)  # 忽略NaN值

                # 将结果添加到DataFrame
                new_row = pd.DataFrame({
                    columns[0]: [admin_code],
                    columns[1]: [admin_code_to_name.get(admin_code, '未知地区')],
                    columns[2]: [quantity_type],
                    columns[3]: [ecosystem_type],
                    columns[4]: [total_value],
                    columns[5]: [mean_value]
                })
                results = pd.concat([results, new_row], ignore_index=True)

    # 根据 folder_path 生成默认的 output_path
    if output_path is None:
        folder_name = os.path.basename(folder_path)
        parent_folder_path = os.path.dirname(folder_path) # 获取上级文件夹路径
        output_path = os.path.join(parent_folder_path, rf'{folder_name}_results.xlsx')

    # 将结果存储到Excel文件
    results.to_excel(output_path, index=False)

    print(f'结果已保存到 {output_path}')

# 示例调用
if __name__ == "__main__":
    tif_folder_path = r'H:/GEP_result/2023/test_Qap'
    analyze_tif_files(tif_folder_path)
